import maya.cmds as mc
import glTools.tools.controlBuilder
import glTools.utils.base
import glTools.utils.channelState

class PivotObject(object):
	
	def __init__(self):
		'''
		Initializer for PivotObject class.
		'''
		# Pivot Control Names
		self.pivotControlName = 'PivotCONTROL'
		self.pivotControlGrp = 'PivotCONTROL_grp'
		# Pivot Locator Names
		self.pivotLocatorName = 'PivotLOCATOR'
		self.pivotLocatorBuf = 'PivotLOCATOR_buf'
		self.pivotLocatorGrp = 'PivotLOCATOR_grp'
		# Channel Arrays
		self.translateChannels = ['tx','ty','tz']
		self.rotateChannels = ['rx','ry','rz']
		self.transformChannels = ['tx','ty','tz','rx','ry','rz']
	
	def create(self,ctrlList=[],pivotObj='',orientPivotControl=False,initializeCleanUpJob=True,constrainTranslate=True,constrainRotate=True,maintainPivotInfo=False):
		'''
		This function will setup a pivotObject control, which will allow you to manipulate
		a group of controls from a common pivot position.
		@param ctrlList: List of controls that will be controled by pivotObject
		@type ctrlList: list
		@param pivotObj: Object whose pivot position will be used for the pivotObject position
		@type pivotObj: str
		@param orientPivotControl: Orient pivotObject to the specified control
		@type orientPivotControl: bool
		@param initializeCleanUpJob: Initialize scriptJob that will cleanup pivotObject control and constraints
		@type initializeCleanUpJob: bool
		@param constrainTranslate: Affect translate channels of controls driven by pivotObject
		@type constrainTranslate: bool
		@param constrainRotate: Affect rotate channels of controls driven by pivotObject
		@type constrainRotate: bool
		@param maintainPivotInfo: Maintain pivotObject information by generating a locator with relevant attribute values
		@type maintainPivotInfo: bool
		'''
		# Check for existing Pivot Object
		if mc.objExists(self.pivotControlName): self.kill()
		
		# Check control list
		for ctrl in ctrlList:
			if not mc.objExists(ctrl): raise UserInputError('Control object '+ctrl+' does not exists!')
		# Check pivot object
		if not mc.objExists(pivotObj): raise UserInputError('Pivot object '+pivotObj+' does not exists!')
		
		# Create pivot control
		pivCtrl = glTools.tools.controlBuilder.create('sphere',self.pivotControlName)['control']
		pivGrp = glTools.utils.base.group(pivCtrl,0)
		# Position pivot control
		mc.delete(mc.pointConstraint(pivotObj,pivGrp,mo=False))
		# Orient pivot control
		if orientPivotControl:
			mc.delete(mc.orientConstraint(pivotObj,pivGrp,mo=False))
		
		# Set channelStates for pivot control and group
		glTools.utils.channelState.ChannelState().setFlags([0,0,0,0,0,0,0,0,0,1],[pivCtrl])
		glTools.utils.channelState.ChannelState().setFlags([2,2,2,2,2,2,2,2,2,1],[pivGrp])
		glTools.utils.channelState.ChannelState().set(1,[pivCtrl,pivGrp])
		
		# Create constraints
		for ctrl in ctrlList:
			# Check translate channels
			skipTranslate = []
			if constrainTranslate:
				for ch in self.translateChannels:
					if not mc.getAttr(ctrl+'.'+ch,se=True):
						skipTranslate.append(ch[1])
			else : skipTranslate = ['x','y','z']
			# Check rotate channels
			skipRotate = []
			if constrainRotate:
				for ch in self.rotateChannels:
					if not mc.getAttr(ctrl+'.'+ch,se=True):
						skipRotate.append(ch[1])
			else : skipRotate = ['x','y','z']
			# Generate parentConstraint
			mc.parentConstraint(pivCtrl,ctrl,skipTranslate=skipTranslate,skipRotate=skipRotate,mo=True)
		
		# Select pivot control
		mc.select(self.pivotControlName)
		
		# Setup kill scriptJob
		if initializeCleanUpJob: self.cleanUpJob(maintainPivotInfo)
	
	def cleanUp(self,maintainPivotInfo=False):
		'''
		This function will remove and cleanup any nodes, controls and constraints generated by PivotObject.create()
		@param maintainPivotInfo: Maintain pivotObject information by generating a locator with relevant attribute values
		@type maintainPivotInfo: bool
		'''
		# Check pivot control exists
		if mc.objExists(self.pivotControlName):
			
			# Check maintain info
			if maintainPivotInfo: self.pivotLocator()
			
			# Get constraint/control lists
			controlList = []
			constraintList = []
			[constraintList.append(conn) for conn in mc.listConnections('PivotCONTROL',s=False,d=True,type='constraint') if not constraintList.count(conn)]
			for constraint in constraintList:
				constraintSlave = mc.listConnections(constraint+'.constraintRotatePivot',s=1,d=0)
				if constraintSlave: controlList.append(constraintSlave[0])
			# Record control transform values
			controlTransforms = {}
			for ctrl in controlList:
				controlTransforms[ctrl] = []
				controlTransforms[ctrl].extend(mc.getAttr(ctrl+'.t')[0])
				controlTransforms[ctrl].extend(mc.getAttr(ctrl+'.r')[0])
			
			# Delete constraints and control
			mc.delete(mc.listConnections('PivotCONTROL',s=False,d=True,type='constraint'))
			mc.delete(self.pivotControlName)
			
			# Reapply control transform values
			transformChannels = self.translateChannels + self.rotateChannels
			for ctrl in controlTransforms.iterkeys():
				for i in range(len(transformChannels)):
					try: mc.setAttr(ctrl+'.'+transformChannels[i],controlTransforms[ctrl][i])
					except: pass
		
		# Delete control group
		if mc.objExists(self.pivotControlGrp):
			mc.delete(self.pivotControlGrp)
	
	def cleanUpJob(self,maintainPivotInfo=False):
		'''
		Initializes a scriptJob that will run a function that will remove and cleanup any nodes, controls and constraints
		generated by PivotObject.create()
		@param maintainPivotInfo: Maintain pivotObject information by generating a locator with relevant attribute values
		@type maintainPivotInfo: bool
		'''
		mm.eval('scriptJob -e "SelectionChanged" "python(\\\"PivotObject().cleanUp('+str(maintainPivotInfo)+')\\\")" -runOnce true')
		
	def pivotLocator(self):
		'''
		Generates a pivotObject locator that will contain all necessary information (via attribute values)
		to recreate a pivotObject.
		'''
		# Check pivot locator group
		if not mc.objExists(self.pivotLocatorGrp):
			mc.group(n=self.pivotLocatorGrp,em=True)
			glTools.utils.channelState.ChannelState().setFlags([2,2,2,2,2,2,2,2,2,1],[self.pivotLocatorGrp])
			glTools.utils.channelState.ChannelState().set(1,[self.pivotLocatorGrp])
		
		# Create pivot locator
		#pivLoc = glTools.tools.controlBuilder.create('locator',self.pivotLocatorName)
		pivLoc = mc.group(n=self.pivotLocatorName,em=True)
		# Create pivot locator annotation
		pivLocAnnot = mc.createNode('annotationShape',p=pivLoc)
		mc.setAttr(pivLocAnnot+'.text',self.pivotLocatorName,type='string')
		mc.setAttr(pivLocAnnot+'.displayArrow',0)
		mc.delete(mc.parentConstraint(self.pivotControlName,pivLoc))
		# Create pivot locator group
		pivLocBuf = mc.group(n=self.pivotLocatorBuf,em=True)
		mc.delete(mc.parentConstraint(self.pivotControlGrp,pivLocBuf))
		# Parent locator to group
		mc.parent(pivLoc,pivLocBuf)
		mc.parent(pivLocBuf,self.pivotLocatorGrp)
		# Set pivot locator channel states
		glTools.utils.channelState.ChannelState().setFlags([2,2,2,2,2,2,2,2,2,1],[pivLoc,pivLocBuf])
		glTools.utils.channelState.ChannelState().set(1,[pivLoc,pivLocBuf])
		
		# Record Control List
		#---------------------
		# Get affected control list
		controlList = []
		constraintList = []
		constraintConnList = mc.listConnections(self.pivotControlName,s=False,d=True,type='constraint')
		[constraintList.append(conn) for conn in constraintConnList if not constraintList.count(conn)]
		constraintList.sort()
		for constraint in constraintList:
			constraintSlave = mc.listConnections(constraint+'.constraintRotatePivot',s=1,d=0)
			if constraintSlave: controlList.append(constraintSlave[0])
		# Add control list values to locator string array attribute
		if len(controlList):
			mc.addAttr(pivLoc,ln='controlList',m=True,dt='string')
			for ctrl in range(len(controlList)):
				mc.setAttr(pivLoc+'.controlList['+str(ctrl)+']',controlList[ctrl],type='string')
		mc.setAttr(pivLoc+'.controlList',l=True)
		# Record Current Frame Number
		#----------------------------
		mc.addAttr(pivLoc,ln='frame',at='float')
		mc.setAttr(pivLoc+'.frame',mc.currentTime(q=True))
		mc.setAttr(pivLoc+'.frame',l=True)
	

